import pandas as pd
import seaborn as sns


level = 'hard'
csv_path = f'dataset/third_time_images/{level}.csv'

info_paths = [
    f'ablation_exps/Surviving/{level}/TranDGN-tran/run1_SEED_0/info.txt',
    f'ablation_exps/Surviving/{level}/TranDGN-tran/run2_SEED_1/info.txt'
]

scores = []
episodes = []
model_types = []
with open(csv_path, 'r') as f:
    lines = f.readlines()[1:]
    for line in lines:
        _id, score, episode, model_type = line.split(',')
        if model_type == 'TranDGN\n':
            scores.append(float(score))
            episodes.append(int(episode))
            model_types.append('TMAC')

for info_path in info_paths:
    with open(info_path, 'r') as f:
        lines = f.readlines()
        for line in lines:
            episode, score = line.split(' ')

            scores.append(float(score) * 10)
            episodes.append(int(episode))
            model_types.append('TMAC w/ MLP')

sns.set_style('whitegrid')
pd_data = pd.DataFrame({'scores': scores, 'types': model_types, 'episodes': episodes})
pd_data.to_csv(f'{level}_with_MLP.csv')

# hue_order = ['TMAC', 'TMAC w/ MLP']
# facetGrid = sns.relplot(data=pd_data, kind='line', x='episodes', y='scores', hue='types', hue_order=hue_order, palette=sns.color_palette(n_colors=len(hue_order)))
# facetGrid.set(xlabel='episodes', ylabel='Reward')
# facetGrid.legend.set_loc('lower right')
# facetGrid.legend.set_title('')
# facetGrid.legend.set_bbox_to_anchor((0.7, 0.1))
# facetGrid.savefig(f"{level}_with_MLP.pdf")